package com.gosling.cloudspace.aop;


import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;

public class MdcHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {
    @Override
    public Callable wrapCallable(Callable callable) {
        return new MdcAwareCallable(callable, MDC.getCopyOfContextMap());
    }

    private class MdcAwareCallable<V> implements Callable {

        private final Callable delegate;

        private final Map contextMap;

        public MdcAwareCallable(Callable callable, Map contextMap) {
            this.delegate = callable;
            this.contextMap = contextMap != null ? contextMap : new HashMap();
        }

        /**
         * 重写调用过程，调用前设置MDC
         * 调用结束后清除
         */
        @Override
        public V call() throws Exception {
            try {
                MDC.setContextMap(contextMap);
                return (V) delegate.call();
            } finally {
                MDC.clear();
            }
        }
    }

}